"""
分析 trace 文件中的 spans 数据。这里是简单计数。
"""

import argparse
import sys
from pathlib import Path

from trace_loader import TraceLoader


class SpanAnalyzer:
    """Span 分析器，用于聚合和分析 span 数据"""
    
    def __init__(self, trace_loader):
        """
        初始化 SpanAnalyzer
        
        Args:
            trace_loader: TraceLoader 实例
        """
        self.trace_loader = trace_loader
        self.in_spans_by_process = trace_loader.in_spans_by_process
        self.out_spans_by_process = trace_loader.out_spans_by_process
        self.all_spans = trace_loader.all_spans
        self.all_processes = trace_loader.all_processes
    
    def get_all_services(self):
        """获取所有服务列表"""
        in_services = set(self.in_spans_by_process.keys())
        out_services = set(self.out_spans_by_process.keys())
        return sorted(in_services | out_services)
    
    def get_service_stats(self, service):
        """
        获取单个服务的统计信息
        
        Args:
            service: 服务名称
            
        Returns:
            dict: 包含 in_spans 和 out_spans 数量的字典
        """
        in_count = len(self.in_spans_by_process.get(service, []))
        out_count = len(self.out_spans_by_process.get(service, []))
        
        return {
            'service': service,
            'in_spans_count': in_count,
            'out_spans_count': out_count,
            'total_spans': in_count + out_count
        }
    
    def get_all_services_stats(self):
        """获取所有服务的统计信息"""
        services = self.get_all_services()
        stats = []
        
        for service in services:
            stats.append(self.get_service_stats(service))
        
        return stats
    
    def print_service_stats(self, service):
        """
        打印单个服务的统计信息
        
        Args:
            service: 服务名称
        """
        stats = self.get_service_stats(service)
        print(f"\n{'='*60}")
        print(f"Service: {stats['service']}")
        print(f"{'-'*60}")
        print(f"  In-Spans:  {stats['in_spans_count']:>6}")
        print(f"  Out-Spans: {stats['out_spans_count']:>6}")
        print(f"  Total:     {stats['total_spans']:>6}")
        print(f"{'='*60}")
    
    def print_all_services_stats(self):
        """打印所有服务的统计信息"""
        services = self.get_all_services()
        total_in = 0
        total_out = 0
        
        print("\n" + "="*80)
        print(f"{'Service Name':<40} {'In-Spans':>12} {'Out-Spans':>12} {'Total':>12}")
        print("="*80)
        
        for service in services:
            stats = self.get_service_stats(service)
            print(f"{stats['service']:<40} {stats['in_spans_count']:>12} "
                  f"{stats['out_spans_count']:>12} {stats['total_spans']:>12}")
            total_in += stats['in_spans_count']
            total_out += stats['out_spans_count']
        
        print("-"*80)
        print(f"{'TOTAL':<40} {total_in:>12} {total_out:>12} {total_in + total_out:>12}")
        print("="*80)
        
        # 打印额外的统计信息
        print(f"\n统计摘要:")
        print(f"  总服务数: {len(services)}")
        print(f"  总 Span 数: {len(self.all_spans)}")
        print(f"  总 Trace 数: {len(self.all_processes)}")
    
    def get_service_endpoints(self, service):
        """
        获取服务的所有端点（endpoint）
        
        Args:
            service: 服务名称
            
        Returns:
            dict: 包含 in_endpoints 和 out_endpoints 的字典
        """
        in_endpoints = set()
        out_endpoints = set()
        
        # 获取 in_spans 的端点
        for span in self.in_spans_by_process.get(service, []):
            ep = span.GetParentEndpoint(self.all_processes, self.all_spans)
            if ep:
                in_endpoints.add(ep)
        
        # 获取 out_spans 的端点
        for span in self.out_spans_by_process.get(service, []):
            ep = span.GetChildEndpoint(self.all_processes, self.all_spans)
            if ep:
                out_endpoints.add(ep)
        
        return {
            'in_endpoints': sorted(in_endpoints),
            'out_endpoints': sorted(out_endpoints)
        }
    
    def print_service_endpoints(self, service):
        """
        打印服务的端点信息
        
        Args:
            service: 服务名称
        """
        endpoints = self.get_service_endpoints(service)
        
        print(f"\n{'='*60}")
        print(f"Service: {service}")
        print(f"{'-'*60}")
        print(f"Incoming Endpoints ({len(endpoints['in_endpoints'])}):")
        for ep in endpoints['in_endpoints']:
            print(f"  - {ep}")
        
        print(f"\nOutgoing Endpoints ({len(endpoints['out_endpoints'])}):")
        for ep in endpoints['out_endpoints']:
            print(f"  - {ep}")
        print(f"{'='*60}")
    
    def export_to_csv(self, output_file):
        """
        导出统计信息到 CSV 文件
        
        Args:
            output_file: 输出文件路径
        """
        import csv
        
        stats = self.get_all_services_stats()
        
        with open(output_file, 'w', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=['service', 'in_spans_count', 'out_spans_count', 'total_spans'])
            writer.writeheader()
            writer.writerows(stats)
        
        print(f"\n统计信息已导出到: {output_file}")


def main():
    parser = argparse.ArgumentParser(description='分析 trace 文件中的 spans 数据')
    parser.add_argument('--relative_path', type=str, required=False, default=None,
                        help='trace 文件的相对路径')
    parser.add_argument('--absolute_path', type=str, required=False, default=None,
                        help='trace 文件的绝对路径')
    parser.add_argument('--compressed', action='store_true',
                        help='目录是否压缩')
    parser.add_argument('--clear_cache', action='store_true',
                        help='清除缓存文件')
    parser.add_argument('--max_traces', type=int, default=None,
                        help='最大加载 trace 数量')
    parser.add_argument('--fix_mode', type=int, default=5, choices=[0, 1, 2, 3, 4, 5, 6, 7],
                        help='修复模式 (0-7)')
    parser.add_argument('--verbose', action='store_true',
                        help='输出详细信息')
    parser.add_argument('--debug', action='store_true',
                        help='调试模式（只加载 50%% 数据）')
    parser.add_argument('--service', type=str, default=None,
                        help='只分析指定的服务')
    parser.add_argument('--show_endpoints', action='store_true',
                        help='显示服务的端点信息')
    parser.add_argument('--export_csv', type=str, default=None,
                        help='导出统计信息到 CSV 文件')
    
    args = parser.parse_args()
    
    # 检查路径参数
    if args.relative_path is None and args.absolute_path is None:
        parser.error("必须指定 --relative_path 或 --absolute_path 之一")
    
    # 确定 trace 目录
    if args.absolute_path:
        trace_dir = args.absolute_path
    else:
        import os
        import helpers.misc as misc
        project_root = misc.get_project_root()
        trace_dir = os.path.join(project_root, args.relative_path)
    
    print(f"正在加载 trace 文件...")
    print(f"  目录: {trace_dir}")
    print(f"  修复模式: {args.fix_mode}")
    print(f"  最大 trace 数: {args.max_traces if args.max_traces else '无限制'}")
    
    # 创建 TraceLoader 并加载数据
    loader = TraceLoader(
        fix_mode=args.fix_mode,
        verbose=args.verbose,
        debug=args.debug
    )
    
    in_spans_by_process, out_spans_by_process, all_spans, all_processes = loader.load_traces(
        directory=trace_dir,
        compressed=args.compressed,
        clear_cache=args.clear_cache,
        max_traces=args.max_traces
    )
    
    print(f"\n加载完成！")
    print(f"  总 Span 数: {len(all_spans)}")
    print(f"  总 Trace 数: {len(all_processes)}")
    print(f"  服务数: {len(set(in_spans_by_process.keys()) | set(out_spans_by_process.keys()))}")
    
    # 创建分析器
    analyzer = SpanAnalyzer(loader)
    
    # 根据参数执行不同的分析
    if args.service:
        # 只分析指定的服务
        analyzer.print_service_stats(args.service)
        if args.show_endpoints:
            analyzer.print_service_endpoints(args.service)
    else:
        # 分析所有服务
        analyzer.print_all_services_stats()
        
        if args.show_endpoints:
            print("\n" + "="*80)
            print("服务端点详情")
            print("="*80)
            for service in analyzer.get_all_services():
                analyzer.print_service_endpoints(service)
    
    # 导出到 CSV
    if args.export_csv:
        analyzer.export_to_csv(args.export_csv)


if __name__ == '__main__':
    main()
